import numpy as np
import pandas as pd
from matplotlib import pyplot as  plt
from sklearn.model_selection import train_test_split
import torch
from torch import nn
from torch.nn import functional as F

device="cuda:0" if torch.cuda.is_available() else "cpu"

"""
1. ✨✨✨ 准备数据✨✨✨ 
   pytorch要求对数据进行打包，打包数据的方法分两步：
    a.重载Dataset类。
    b.DataLoader 数据加载器。
"""
from torch.utils.data import Dataset
from torch.utils.data import DataLoader


df = pd.read_excel('class/the boston_house_prices.xlsx', engine='openpyxl')  
X = df.drop('MEDV', axis=1).values
y = df['MEDV'].values


#切分数据
X_train,X_test,y_train,y_test = train_test_split(X,y,test_size=0.2,random_state=1)
y_train = y_train.reshape((-1,1))
y_test = y_test.reshape((-1,1))


#进行数据scalar
mean_ = X_train.mean(axis=0)
std_ = X_train.std(axis=0)
#print(mean_.shape)
#print(X_test.shape)
X_train= (X_train-mean_)/std_

X_test = (X_test-mean_)/std_


#重载dataset类
class MyDataset(Dataset):
   def __init__(self,X,y):
       self.X = X
       self.y = y
   #重载__getitem__方法。
   def __getitem__(self, idx):
       return torch.tensor(data=self.X[idx]).float(),torch.tensor(data=self.y[idx]).float()
   def __len__(self):
       return len(self.X)

#以下：定义训练集的数据加载器
train_dataset = MyDataset(X=X_train,y=y_train)

#测试代码
#print(train_dataset[0]) #根据idex输出对应的训练样本特征
#封装为dataloader。Shuffle是指打乱顺序。droplast是把不足的那一批给干掉了。
train_dataloader = DataLoader(dataset=train_dataset,batch_size=32,shuffle=True,drop_last=True)
#for X,y in train_dataloader:
#   print(X.shape,y.shape)

#以下定义测试集的数据加载器
test_dataset = MyDataset(X=X_test,y=y_test)

test_dataloader = DataLoader(dataset=test_dataset,batch_size=32,shuffle=False,drop_last=False)

"""
2.✨✨✨ 定义模型✨✨✨ 
"""
class LinearRegression(nn.Module):
   def __init__(self):
       super(LinearRegression,self).__init__() #一个子类继承一个父类，所以子类在构建的时候，要先构建一个父类。此行代码写法固定。
       self.linear1 = nn.Linear(in_features=13,out_features=5)
       self.linear2 = nn.Linear(in_features=5,out_features=1)

   def forward(self,x):
       """
       定义正向传播的过程。
       当参数已经训练好了以后，正向传播就是正向推理即predict。
       x的形状必须是[b,in_featurs],out的形状是[b,1]。b是batchsize。
       """
       out = self.linear1(x) #批量的计算。hidden1
       out = F.relu(out)
       out = self.linear2(out)  # 批量的计算。
       return out
lr = LinearRegression()
#此处可以统一将模型的参数转移到指定的设备上：
lr.to(device)

# 调试代码。参数定义和业务逻辑都定义在层里面了。我们现在可以直接来看看模型参数：
#print(lr.parameters())
#for ele in lr.parameters():
#   print(ele)

"""
3.✨✨✨ 模型训练✨✨✨ 
模型训练包括：
指定损失函数；
指定优化器；
进行梯度下降
"""

#定义损失函数
loss_fn = nn.MSELoss()
#定义优化器（以下是指把参数lr.parameters托管给优化器SGD）
optimizer = torch.optim.SGD(params=lr.parameters(),lr=1e-2)
#训练过程监控函数
def get_loss(dataloader=train_dataloader,model=lr,loss_fn=loss_fn):
   with torch.no_grad():
       model.eval() #评估模式。
       losses = []
       for X,y in dataloader:
           #此处可以将数据指定设备。配合前面lr.to（device=device)使用。
           X = X.to(device=device)
           y = y.to(device=device)
           y_pred = model(X) #批量计算。"model的参数形状和X要能对上口型。才能参与矩阵运算。" 。
           loss = loss_fn(y_pred,y) #结果也在对应的device上。
           losses.append(loss.item())
       return np.array(losses).mean()

#利用梯度下降开始训练
def train(model,dataloader, optimizer=optimizer,loss_fn=loss_fn,epochs=20):
   model.train()
   for epoch in range(1,epochs+1):
       #每一轮进行训练
       for X,y in dataloader:
           X = X.to(device=device)
           y = y.to(device=device)
           #正向传播：
           y_pred = model(X) #这个一调用实际上经过了forward过程。对于多层网络的情况，这样关系是函数嵌套：f3(f2(f1（x）))
           #计算损失：
           loss = loss_fn(y_pred,y) #loss = f(w,b) 。loss是单值函数。这里就是一个单值函数对一堆变量求偏导。

#清空梯度：
           optimizer.zero_grad()


           #反向传播求导：
           loss.backward() #loss = f(w,b) 。loss是单值函数。这里就是一个单值函数对一堆变量求偏导。
           #优化 step的意思就是减去一次梯度：
           optimizer.step() #更新参数：
       #以上操作更新了模型。而且以上每走一步就更新了一次。
       print(f"当前是第{epoch}轮，训练集的误差为：{get_loss(dataloader=train_dataloader)},测试集的误差为：{get_loss(dataloader=test_dataloader)}")
train(lr, train_dataloader, optimizer, loss_fn)
#测试代码
#print(lr.state_dict())
#print(len(lr.state_dict()["linear.weight"][0]))

#5.✨✨✨ 模型预测✨✨✨ 
def predict(model,X):
   model.eval()
   #其实所谓的预测就是正向传播一下，但是不用构建计算图。 with torch.no_grad().
   with torch.no_grad():
       X_tensor = X.to(device)
       y_pred = model(X_tensor)
   return y_pred